# MIT License
#
# Copyright (c) 2023 Dechin CHEN
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import os
os.environ['GLOG_v'] = '4'
from quaternion import Quaternion
from quaternion.functions import quaternion_multiply, quaternion_inverse, hamiltonian_product, quaternion_diff
from quaternion.functions import batch_quaternion_multiply, batch_hamiltonian_product, batch_quaternion_diff
import numpy as np
import mindspore as ms
from mindspore import Tensor

if __name__ == '__main__':
    q1 = Tensor([1, 2, 3, 4], ms.float32)
    q1_quaternion = Quaternion(q1)
    q2 = Tensor([3, 1, 0, 5], ms.float32)
    q2_quaternion = Quaternion(q2)
    vector = Tensor([1, 2, 3], ms.float32)
    vector_quaternion = Quaternion(vector)
    hamiltonian = Tensor([0.5, 0.5, 0.5, 0.5], ms.float32)
    hamiltonian_quaternion = Quaternion(hamiltonian)
    src_vector = Tensor([1, 0, 0], ms.float32)
    src_vector_quaternion = Quaternion(src_vector)
    dst_vector1 = Tensor([0, 0, 1], ms.float32)
    dst_vector1_quaternion = Quaternion(dst_vector1)
    dst_vector2 = Tensor([np.sqrt(3) / 3, np.sqrt(3) / 3, np.sqrt(3) / 3], ms.float32)
    dst_vector2_quaternion = Quaternion(dst_vector2)

    print ('The conjugate of q1 is: {}'.format(q1_quaternion.conjugate()))
    assert np.allclose(q1_quaternion.conjugate().to_tensor().asnumpy(), np.array([[1., -2., -3., -4.]], np.float32))
    print ('The conjugate of hamiltonian is: {}'.format(quaternion_inverse(hamiltonian)))
    assert np.allclose(quaternion_inverse(hamiltonian).asnumpy(), np.array([[0.5, -0.5, -0.5, -0.5]], np.float32))
    print ('The sum of q1 and q2 is: {}'.format(q1_quaternion + q2_quaternion))
    assert np.allclose((q1_quaternion + q2_quaternion).to_tensor().asnumpy(), np.array([[4., 3., 3., 9.]], np.float32))
    # The sum of q1 and q2 is: [[4. 3. 3. 9.]]
    print('The sub of q2 is: {}'.format(q1_quaternion - q2_quaternion))
    assert np.allclose((q1_quaternion - q2_quaternion).to_tensor().asnumpy(), np.array([[-2., 1., 3., -1.]], np.float32))
    # The sub of q2 is: [[-2.  1.  3. -1.]]
    print('The product value of q1 and -1 is: {}'.format(q1_quaternion * (-1)))
    assert np.allclose((q1_quaternion * (-1)).to_tensor().asnumpy(), np.array([[-1., -2., -3., -4.]], np.float32))
    # The product value of q1 and -1 is: [[-1. -2. -3. -4.]]
    print('The product value of -1 and q1 is: {}'.format(-1 * q1_quaternion))
    assert np.allclose((-1 * q1_quaternion).to_tensor().asnumpy(), np.array([[-1., -2., -3., -4.]], np.float32))
    # The product value of -1 and q1 is: [[-1. -2. -3. -4.]]
    print ('The product of q1 and q2 is: {}'.format(q1_quaternion * q2_quaternion))
    assert np.allclose((q1_quaternion * q2_quaternion).to_tensor().asnumpy(),
                       np.array([[-19., 22., 3., 14.]], np.float32))
    assert np.allclose(quaternion_multiply(q1, q2).asnumpy(), np.array([[-19., 22., 3., 14.]], np.float32))
    # The product of q1 and q2 is: [[-19.  22.   3.  14.]]
    print('The product of q2 and q1 is: {}'.format(q2_quaternion * q1_quaternion))
    assert np.allclose((q2_quaternion * q1_quaternion).to_tensor().asnumpy(),
                       np.array([[-19., -8., 15., 20.]], np.float32))
    assert np.allclose(quaternion_multiply(q2, q1).asnumpy(), np.array([[-19., -8., 15., 20.]], np.float32))
    # The product of q2 and q1 is: [[-19.  -8.  15.  20.]]
    print ('The shape of q1 is: {}'.format(len(q1_quaternion)))
    assert len(q1_quaternion) == 1
    # The shape of q1 is: 1
    print ('The absolute of q2 is: {}'.format(abs(q2_quaternion)))
    assert np.allclose(abs(q2_quaternion).asnumpy(), np.array([5.916079], np.float32))
    # The absolute of q2 is: [5.916079]
    print('The value of 1/q1 is: {}'.format(1/q1_quaternion))
    assert np.allclose((1/q1_quaternion).to_tensor().asnumpy(),
                       np.array([[0.03333333, -0.06666666, -0.09999999, -0.13333333]], np.float32))
    # The value of 1/q1 is: [[ 0.03333333 -0.06666666 -0.09999999 -0.13333333]]
    print ('The value of q1/2 is: {}'.format(q1_quaternion/2))
    assert np.allclose((q1_quaternion/2).to_tensor().asnumpy(),
                       np.array([[0.5, 1.,  1.5, 2. ]], np.float32))
    # The value of q1/2 is: [[0.5 1.  1.5 2. ]]
    print ('The Hamiltonian product of hamiltonian | vector > is: {}'.format(hamiltonian_quaternion | vector))
    assert np.allclose((hamiltonian_quaternion | vector).asnumpy(), np.array([[3., 1., 2.]], np.float32))
    assert np.allclose(hamiltonian_product(hamiltonian, vector).asnumpy(), np.array([[0., 3., 1., 2.]], np.float32))
    # The Hamiltonian product of hamiltonian | vector > is: [[3. 1. 2.]]
    vector_transform = hamiltonian_quaternion.conjugate() | (hamiltonian_quaternion | vector)
    print ('The Hamiltonian product of hamiltonian^(-1) | hamiltonian | vector > is: {}'.format(vector_transform))
    assert np.allclose(vector_transform.asnumpy(), np.array([[1., 2., 3.]], np.float32))
    # The Hamiltonian product of hamiltonian^(-1) | hamiltonian | vector > is: [[1. 2. 3.]]
    transform_quaternion = src_vector_quaternion >> dst_vector1_quaternion
    print ('The quaternion transform from vector to target_vector is: {}'.format(transform_quaternion))
    assert np.allclose(transform_quaternion.to_tensor().asnumpy(),
                       np.array([[0.70710677,  0.,         -0.7071068,   0.]], np.float32))
    assert np.allclose(quaternion_diff(src_vector, dst_vector1).asnumpy(),
                       np.array([[0.70710677,  0.,         -0.7071068,   0.]], np.float32))
    # The quaternion transform from vector to target_vector is: [[ 0.70710677  0.         -0.7071068   0.        ]]
    print ('The transform result is: {}'.format(transform_quaternion | src_vector))
    assert np.allclose((transform_quaternion | src_vector).asnumpy(), np.array([[0., 0., 1.]], np.float32))
    # The transform result is: [[0.         0.         0.99999994]]
    transform_quaternion = src_vector_quaternion >> dst_vector2_quaternion
    print('The quaternion transform from vector to target_vector is: {}'.format(transform_quaternion))
    assert np.allclose(transform_quaternion.to_tensor().asnumpy(),
                       np.array([[0.88807386,  0.,         -0.32505763,  0.32505763]], np.float32))
    assert np.allclose(quaternion_diff(src_vector, dst_vector2).asnumpy(),
                       np.array([[0.88807386,  0.,         -0.32505763,  0.32505763]], np.float32))
    # The quaternion transform from vector to target_vector is: [[ 0.88807386  0.         -0.32505763  0.32505763]]
    print('The transform result is: {}'.format(transform_quaternion | src_vector))
    assert np.allclose((transform_quaternion | src_vector).asnumpy(),
                       np.array([[0.57735026, 0.5773504,  0.5773504]], np.float32))
    # The transform result is: [[0.57735026 0.5773504  0.5773504 ]]

    np.random.seed(1)
    batch_src_vector = np.random.random((10, 3))
    batch_src_vector /= np.linalg.norm(batch_src_vector, axis=-1)[:, None]
    batch_src_vector = Tensor(batch_src_vector, dtype=ms.float32)
    batch_src_vector_quaternion = Quaternion(batch_src_vector)
    batch_dst_vector = np.random.random((10, 3))
    batch_dst_vector /= np.linalg.norm(batch_dst_vector, axis=-1)[:, None]
    batch_dst_vector = Tensor(batch_dst_vector, dtype=ms.float32)
    batch_dst_vector_quaternion = Quaternion(batch_dst_vector)
    batch_transform_quaternion = batch_src_vector_quaternion >> batch_dst_vector_quaternion
    batch_transform_vector = batch_transform_quaternion | batch_src_vector
    answer = np.allclose(batch_dst_vector.asnumpy(), batch_transform_vector.asnumpy())
    assert answer
    print ('Is the transformed batch vector equal to batch src vector? {}'.format(answer))

    # Test batch operations.
    q1 = q1[None, :]
    q2 = q2[None, :]
    mul1 = quaternion_multiply(q1, q2)[None, :]
    ham1 = hamiltonian_product(q1, q2)[None, :]
    dif1 = quaternion_diff(q1, q2)[None, :]
    q1 = q1[None, :]
    q2 = q2[None, :]
    mul2 = batch_quaternion_multiply(q1, q2)
    ham2 = batch_hamiltonian_product(q1, q2)
    dif2 = batch_quaternion_diff(q1, q2)
    assert np.allclose(mul1.asnumpy(), mul2.asnumpy())
    assert np.allclose(ham1.asnumpy(), ham2.asnumpy())
    assert np.allclose(dif1.asnumpy(), dif2.asnumpy())
